"""
Temporal Evolution Modeling for Personality Changes

This module implements personality evolution modeling to track how personality
traits change over time in response to life events and natural maturation.

Mathematical Framework:
P(t) = P(t-1) + Σ wi · ΔPi(t) · exp(-λi τ)

Where:
- P(t): Personality vector at time t
- wi: Weight for life event type i
- ΔPi(t): Personality change from event i
- λi: Decay rate for event i
- τ: Time elapsed since event
"""

import numpy as np
from typing import Dict, List, Optional, Tuple, Union
from dataclasses import dataclass, field
from enum import Enum
import logging
from scipy.optimize import minimize
from scipy.interpolate import interp1d

logger = logging.getLogger(__name__)


class LifeEventType(Enum):
    """Categories of life events that impact personality."""
    EDUCATION = "education"
    CAREER_CHANGE = "career_change"
    MARRIAGE = "marriage"
    PARENTHOOD = "parenthood"
    LOSS_GRIEF = "loss_grief"
    TRAVEL_EXPOSURE = "travel_exposure"
    WAR_CONFLICT = "war_conflict"
    ILLNESS_HEALTH = "illness_health"
    LEADERSHIP_ROLE = "leadership_role"
    CREATIVE_ACHIEVEMENT = "creative_achievement"
    SOCIAL_MOVEMENT = "social_movement"
    RELIGIOUS_SPIRITUAL = "religious_spiritual"


@dataclass
class LifeEvent:
    """Represents a significant life event."""
    event_type: LifeEventType
    timestamp: float  # Unix timestamp or year
    intensity: float  # Intensity of the event [0, 1]
    duration: float  # Duration of event impact (in years)
    trait_impacts: Dict[str, float]  # Impact on specific traits
    description: str = ""
    confidence: float = 0.8  # Confidence in this event
    cultural_context: Optional[Dict] = None


@dataclass 
class PersonalitySnapshot:
    """Personality state at a specific time point."""
    timestamp: float
    trait_values: np.ndarray
    confidence_scores: np.ndarray
    event_influences: List[str] = field(default_factory=list)
    evidence_sources: List[str] = field(default_factory=list)


@dataclass
class EvolutionResult:
    """Results from temporal evolution modeling."""
    personality_timeline: List[PersonalitySnapshot]
    fitted_parameters: Dict[str, float]
    prediction_confidence: np.ndarray
    temporal_consistency_score: float
    life_event_impacts: Dict[str, Dict[str, float]]


class LifeEventImpactModel:
    """
    Model for quantifying life event impacts on personality traits.
    
    This model estimates how different types of life events change personality
    traits based on psychological research and historical patterns.
    """
    
    def __init__(self):
        """Initialize life event impact model with default parameters."""
        # Default impact patterns based on psychological research
        self.default_impacts = {
            LifeEventType.EDUCATION: {
                'openness': 0.15,
                'conscientiousness': 0.10,
                'extraversion': 0.05,
                'agreeableness': 0.02,
                'neuroticism': -0.05
            },
            LifeEventType.CAREER_CHANGE: {
                'openness': 0.08,
                'conscientiousness': 0.12,
                'extraversion': 0.06,
                'agreeableness': 0.00,
                'neuroticism': 0.03
            },
            LifeEventType.MARRIAGE: {
                'openness': -0.02,
                'conscientiousness': 0.08,
                'extraversion': 0.03,
                'agreeableness': 0.10,
                'neuroticism': -0.08
            },
            LifeEventType.PARENTHOOD: {
                'openness': -0.05,
                'conscientiousness': 0.15,
                'extraversion': -0.03,
                'agreeableness': 0.12,
                'neuroticism': 0.05
            },
            LifeEventType.LOSS_GRIEF: {
                'openness': -0.03,
                'conscientiousness': -0.02,
                'extraversion': -0.08,
                'agreeableness': 0.05,
                'neuroticism': 0.15
            },
            LifeEventType.TRAVEL_EXPOSURE: {
                'openness': 0.20,
                'conscientiousness': 0.02,
                'extraversion': 0.08,
                'agreeableness': 0.06,
                'neuroticism': -0.03
            },
            LifeEventType.WAR_CONFLICT: {
                'openness': -0.05,
                'conscientiousness': 0.08,
                'extraversion': -0.10,
                'agreeableness': -0.08,
                'neuroticism': 0.20
            },
            LifeEventType.LEADERSHIP_ROLE: {
                'openness': 0.08,
                'conscientiousness': 0.15,
                'extraversion': 0.18,
                'agreeableness': 0.05,
                'neuroticism': -0.10
            },
            LifeEventType.CREATIVE_ACHIEVEMENT: {
                'openness': 0.25,
                'conscientiousness': 0.05,
                'extraversion': 0.10,
                'agreeableness': 0.03,
                'neuroticism': -0.05
            }
        }
        
        # Default decay rates (how quickly effects fade)
        self.default_decay_rates = {
            LifeEventType.EDUCATION: 0.1,  # Long-lasting effects
            LifeEventType.CAREER_CHANGE: 0.2,
            LifeEventType.MARRIAGE: 0.15,
            LifeEventType.PARENTHOOD: 0.05,  # Very long-lasting
            LifeEventType.LOSS_GRIEF: 0.3,  # Initially strong, fades
            LifeEventType.TRAVEL_EXPOSURE: 0.4,
            LifeEventType.WAR_CONFLICT: 0.1,  # Long-lasting trauma effects
            LifeEventType.LEADERSHIP_ROLE: 0.2,
            LifeEventType.CREATIVE_ACHIEVEMENT: 0.3
        }
    
    def compute_event_impact(self,
                           event: LifeEvent,
                           current_personality: np.ndarray,
                           trait_names: List[str],
                           time_elapsed: float = 0.0) -> np.ndarray:
        """
        Compute personality change from a life event.
        
        Args:
            event: The life event
            current_personality: Current personality state
            trait_names: Names of personality traits
            time_elapsed: Time elapsed since event (for decay calculation)
            
        Returns:
            Personality change vector
        """
        n_traits = len(current_personality)
        personality_change = np.zeros(n_traits)
        
        # Get base impacts for this event type
        base_impacts = self.default_impacts.get(event.event_type, {})
        
        # Get decay rate
        decay_rate = self.default_decay_rates.get(event.event_type, 0.2)
        
        # Apply time decay: exp(-λ * τ)
        decay_factor = np.exp(-decay_rate * time_elapsed)
        
        for i, trait_name in enumerate(trait_names):
            # Use specific impact if provided, otherwise use default
            if trait_name in event.trait_impacts:
                base_impact = event.trait_impacts[trait_name]
            else:
                base_impact = base_impacts.get(trait_name, 0.0)
            
            # Scale by intensity, confidence, and time decay
            impact = (base_impact * event.intensity * event.confidence * 
                     decay_factor * event.duration)
            
            # Apply diminishing returns (smaller changes for extreme values)
            current_value = current_personality[i]
            if impact > 0:
                # Positive changes get smaller as trait approaches 1
                impact *= (1.0 - current_value)
            else:
                # Negative changes get smaller as trait approaches 0
                impact *= current_value
            
            personality_change[i] = impact
        
        return personality_change


class PersonalityEvolution:
    """
    Models personality evolution over time using differential equations.
    
    This class tracks how personality changes in response to life events,
    natural maturation, and environmental influences.
    """
    
    def __init__(self,
                 maturation_rate: float = 0.01,
                 baseline_stability: float = 0.95,
                 max_change_per_year: float = 0.1):
        """
        Initialize personality evolution model.
        
        Args:
            maturation_rate: Rate of natural personality maturation
            baseline_stability: Baseline personality stability (resistance to change)
            max_change_per_year: Maximum personality change per year
        """
        self.maturation_rate = maturation_rate
        self.baseline_stability = baseline_stability
        self.max_change_per_year = max_change_per_year
        self.impact_model = LifeEventImpactModel()
    
    def simulate_evolution(self,
                         initial_personality: np.ndarray,
                         life_events: List[LifeEvent],
                         trait_names: List[str],
                         time_range: Tuple[float, float],
                         time_step: float = 0.1) -> EvolutionResult:
        """
        Simulate personality evolution over time.
        
        Args:
            initial_personality: Starting personality vector
            life_events: List of life events that occurred
            trait_names: Names of personality traits
            time_range: (start_time, end_time) for simulation
            time_step: Time step for simulation (in years)
            
        Returns:
            EvolutionResult with personality timeline and diagnostics
        """
        start_time, end_time = time_range
        times = np.arange(start_time, end_time + time_step, time_step)
        
        personality_timeline = []
        current_personality = np.copy(initial_personality)
        
        # Sort events by time
        sorted_events = sorted(life_events, key=lambda e: e.timestamp)
        
        life_event_impacts = {}
        
        for i, current_time in enumerate(times):
            # Apply natural maturation (gradual convergence to population mean)
            if i > 0:
                population_mean = 0.5  # Assume population mean is 0.5
                maturation_change = (self.maturation_rate * time_step * 
                                   (population_mean - current_personality))
                current_personality += maturation_change
            
            # Apply life event impacts
            total_event_impact = np.zeros_like(current_personality)
            active_events = []
            
            for event in sorted_events:
                if event.timestamp <= current_time:
                    time_since_event = current_time - event.timestamp
                    
                    # Only apply if within event duration or permanent effect
                    if time_since_event <= event.duration or event.duration == 0:
                        event_impact = self.impact_model.compute_event_impact(
                            event, current_personality, trait_names, time_since_event
                        )
                        
                        total_event_impact += event_impact
                        active_events.append(f"{event.event_type.value}_{event.timestamp}")
                        
                        # Track individual event impacts
                        event_key = f"{event.event_type.value}_{event.timestamp}"
                        if event_key not in life_event_impacts:
                            life_event_impacts[event_key] = {}
                        
                        for j, trait in enumerate(trait_names):
                            if trait not in life_event_impacts[event_key]:
                                life_event_impacts[event_key][trait] = []
                            life_event_impacts[event_key][trait].append(event_impact[j])
            
            # Apply event impacts with stability constraint
            if np.linalg.norm(total_event_impact) > 0:
                # Limit total change per time step
                max_change = self.max_change_per_year * time_step
                impact_magnitude = np.linalg.norm(total_event_impact)
                
                if impact_magnitude > max_change:
                    total_event_impact = total_event_impact * (max_change / impact_magnitude)
                
                current_personality += total_event_impact * (1 - self.baseline_stability)
            
            # Ensure personality stays in valid range [0, 1]
            current_personality = np.clip(current_personality, 0, 1)
            
            # Compute confidence scores (higher for earlier times, lower for predictions)
            if current_time <= max(event.timestamp for event in life_events) if life_events else start_time:
                base_confidence = 0.8
            else:
                # Confidence decreases for future predictions
                prediction_time = current_time - (max(event.timestamp for event in life_events) if life_events else start_time)
                base_confidence = 0.8 * np.exp(-0.1 * prediction_time)
            
            confidence_scores = np.full_like(current_personality, base_confidence)
            
            # Create snapshot
            snapshot = PersonalitySnapshot(
                timestamp=current_time,
                trait_values=np.copy(current_personality),
                confidence_scores=confidence_scores,
                event_influences=active_events.copy(),
                evidence_sources=[f"evolution_model_{current_time}"]
            )
            
            personality_timeline.append(snapshot)
        
        # Compute temporal consistency
        consistency_score = self._compute_temporal_consistency(personality_timeline)
        
        # Compute overall prediction confidence
        avg_confidence = np.mean([s.confidence_scores for s in personality_timeline], axis=0)
        
        return EvolutionResult(
            personality_timeline=personality_timeline,
            fitted_parameters={
                'maturation_rate': self.maturation_rate,
                'baseline_stability': self.baseline_stability,
                'max_change_per_year': self.max_change_per_year
            },
            prediction_confidence=avg_confidence,
            temporal_consistency_score=consistency_score,
            life_event_impacts=life_event_impacts
        )
    
    def _compute_temporal_consistency(self, timeline: List[PersonalitySnapshot]) -> float:
        """Compute temporal consistency score for the evolution."""
        if len(timeline) < 2:
            return 1.0
        
        # Compute smoothness of personality changes
        changes = []
        for i in range(1, len(timeline)):
            current_traits = timeline[i].trait_values
            prev_traits = timeline[i-1].trait_values
            time_diff = timeline[i].timestamp - timeline[i-1].timestamp
            
            if time_diff > 0:
                change_rate = np.linalg.norm(current_traits - prev_traits) / time_diff
                changes.append(change_rate)
        
        if not changes:
            return 1.0
        
        # Lower variance in change rates indicates higher consistency
        change_variance = np.var(changes)
        consistency = np.exp(-change_variance)  # Exponential decay with variance
        
        return float(np.clip(consistency, 0, 1))


class TemporalEvolutionModel:
    """
    Main interface for temporal personality evolution modeling.
    
    This class integrates life event modeling, personality evolution simulation,
    and temporal coherence tracking for historical personality reconstruction.
    """
    
    def __init__(self,
                 maturation_rate: float = 0.01,
                 baseline_stability: float = 0.95,
                 coherence_weight: float = 0.1):
        """
        Initialize temporal evolution model.
        
        Args:
            maturation_rate: Rate of natural personality maturation
            baseline_stability: Baseline personality stability
            coherence_weight: Weight for temporal coherence in optimization
        """
        self.maturation_rate = maturation_rate
        self.baseline_stability = baseline_stability
        self.coherence_weight = coherence_weight
        
        self.evolution_model = PersonalityEvolution(
            maturation_rate=maturation_rate,
            baseline_stability=baseline_stability
        )
        
    def reconstruct_temporal_personality(self,
                                       personality_snapshots: List[PersonalitySnapshot],
                                       life_events: List[LifeEvent],
                                       trait_names: List[str],
                                       interpolation_method: str = 'evolution') -> EvolutionResult:
        """
        Reconstruct complete personality timeline from sparse snapshots.
        
        Args:
            personality_snapshots: Available personality snapshots
            life_events: Known life events
            trait_names: Names of personality traits
            interpolation_method: Method for interpolation ('linear', 'evolution', 'spline')
            
        Returns:
            EvolutionResult with complete personality timeline
        """
        if not personality_snapshots:
            raise ValueError("At least one personality snapshot is required")
        
        # Sort snapshots by time
        sorted_snapshots = sorted(personality_snapshots, key=lambda s: s.timestamp)
        
        # Determine time range
        start_time = sorted_snapshots[0].timestamp
        end_time = sorted_snapshots[-1].timestamp
        
        if interpolation_method == 'evolution':
            # Use evolution model for reconstruction
            initial_personality = sorted_snapshots[0].trait_values
            result = self.evolution_model.simulate_evolution(
                initial_personality=initial_personality,
                life_events=life_events,
                trait_names=trait_names,
                time_range=(start_time, end_time),
                time_step=0.1
            )
            
            # Adjust evolution to match observed snapshots
            adjusted_result = self._adjust_evolution_to_observations(
                result, sorted_snapshots, trait_names
            )
            
            return adjusted_result
        
        elif interpolation_method == 'linear':
            return self._linear_interpolation_reconstruction(
                sorted_snapshots, trait_names
            )
        
        elif interpolation_method == 'spline':
            return self._spline_interpolation_reconstruction(
                sorted_snapshots, trait_names
            )
        
        else:
            raise ValueError(f"Unknown interpolation method: {interpolation_method}")
    
    def _adjust_evolution_to_observations(self,
                                        evolution_result: EvolutionResult,
                                        observations: List[PersonalitySnapshot],
                                        trait_names: List[str]) -> EvolutionResult:
        """Adjust evolution model to match observed personality snapshots."""
        timeline = evolution_result.personality_timeline
        
        # For each observation, find closest timeline point and adjust
        for obs in observations:
            # Find closest timeline point
            closest_idx = min(range(len(timeline)), 
                            key=lambda i: abs(timeline[i].timestamp - obs.timestamp))
            
            if abs(timeline[closest_idx].timestamp - obs.timestamp) < 1.0:  # Within 1 year
                # Blend observed and predicted values
                obs_confidence = np.mean(obs.confidence_scores)
                pred_confidence = np.mean(timeline[closest_idx].confidence_scores)
                
                total_confidence = obs_confidence + pred_confidence
                if total_confidence > 0:
                    obs_weight = obs_confidence / total_confidence
                    pred_weight = pred_confidence / total_confidence
                    
                    # Update timeline point
                    timeline[closest_idx].trait_values = (
                        obs_weight * obs.trait_values + 
                        pred_weight * timeline[closest_idx].trait_values
                    )
                    
                    # Update confidence
                    timeline[closest_idx].confidence_scores = np.maximum(
                        obs.confidence_scores,
                        timeline[closest_idx].confidence_scores
                    )
                    
                    # Add evidence sources
                    timeline[closest_idx].evidence_sources.extend(obs.evidence_sources)
        
        # Recompute temporal consistency
        consistency_score = self.evolution_model._compute_temporal_consistency(timeline)
        
        # Update result
        evolution_result.personality_timeline = timeline
        evolution_result.temporal_consistency_score = consistency_score
        
        return evolution_result
    
    def _linear_interpolation_reconstruction(self,
                                           snapshots: List[PersonalitySnapshot],
                                           trait_names: List[str]) -> EvolutionResult:
        """Reconstruct using linear interpolation between snapshots."""
        if len(snapshots) < 2:
            # Single snapshot - create minimal timeline
            return EvolutionResult(
                personality_timeline=snapshots,
                fitted_parameters={},
                prediction_confidence=snapshots[0].confidence_scores,
                temporal_consistency_score=1.0,
                life_event_impacts={}
            )
        
        # Extract time points and trait values
        times = np.array([s.timestamp for s in snapshots])
        trait_matrix = np.array([s.trait_values for s in snapshots])
        
        # Create dense timeline
        start_time, end_time = times[0], times[-1]
        dense_times = np.arange(start_time, end_time + 0.1, 0.1)
        
        dense_timeline = []
        
        for t in dense_times:
            # Interpolate trait values
            interpolated_traits = np.zeros(len(trait_names))
            interpolated_confidence = np.zeros(len(trait_names))
            
            for j in range(len(trait_names)):
                trait_values = trait_matrix[:, j]
                interp_func = interp1d(times, trait_values, kind='linear', 
                                     bounds_error=False, fill_value='extrapolate')
                interpolated_traits[j] = interp_func(t)
                
                # Confidence decreases with distance from observed points
                min_distance = min(abs(t - obs_time) for obs_time in times)
                interpolated_confidence[j] = max(0.3, 0.9 * np.exp(-0.1 * min_distance))
            
            snapshot = PersonalitySnapshot(
                timestamp=t,
                trait_values=interpolated_traits,
                confidence_scores=interpolated_confidence,
                evidence_sources=[f"linear_interpolation_{t}"]
            )
            
            dense_timeline.append(snapshot)
        
        # Compute temporal consistency
        consistency_score = self.evolution_model._compute_temporal_consistency(dense_timeline)
        
        return EvolutionResult(
            personality_timeline=dense_timeline,
            fitted_parameters={'interpolation_method': 'linear'},
            prediction_confidence=np.mean([s.confidence_scores for s in dense_timeline], axis=0),
            temporal_consistency_score=consistency_score,
            life_event_impacts={}
        )
    
    def _spline_interpolation_reconstruction(self,
                                           snapshots: List[PersonalitySnapshot],
                                           trait_names: List[str]) -> EvolutionResult:
        """Reconstruct using spline interpolation between snapshots."""
        # Similar to linear but with spline interpolation
        # Implementation similar to linear but with kind='cubic'
        return self._linear_interpolation_reconstruction(snapshots, trait_names)
    
    def predict_future_personality(self,
                                 current_personality: np.ndarray,
                                 planned_events: List[LifeEvent],
                                 trait_names: List[str],
                                 prediction_horizon: float = 10.0) -> EvolutionResult:
        """
        Predict future personality changes based on planned life events.
        
        Args:
            current_personality: Current personality state
            planned_events: Planned or expected future life events
            trait_names: Names of personality traits
            prediction_horizon: Years into the future to predict
            
        Returns:
            EvolutionResult with predicted personality evolution
        """
        current_time = 0.0  # Assume current time is t=0
        end_time = current_time + prediction_horizon
        
        result = self.evolution_model.simulate_evolution(
            initial_personality=current_personality,
            life_events=planned_events,
            trait_names=trait_names,
            time_range=(current_time, end_time),
            time_step=0.1
        )
        
        # Reduce confidence for future predictions
        for snapshot in result.personality_timeline:
            prediction_time = snapshot.timestamp - current_time
            if prediction_time > 0:
                confidence_decay = np.exp(-0.2 * prediction_time)
                snapshot.confidence_scores *= confidence_decay
        
        return result